import torch
from torch import nn as nn


def orthogonal_init(layer, gain=1.0):
    nn.init.orthogonal_(layer.weight, gain=gain)
    nn.init.constant_(layer.bias, 0)


class Actor(nn.Module):
    def __init__(self, state_dim, action_dim, max_action, hyper):
        super().__init__()
        hidden_dim = hyper['hidden_dim']
        use_orthogonal_init = hyper['use_orthogonal_init']
        use_tanh = hyper['use_tanh']
        self.max_action = max_action
        self.l1 = nn.Linear(state_dim, hidden_dim)
        self.l2 = nn.Linear(hidden_dim, hidden_dim)
        self.mean_layer = nn.Linear(hidden_dim, action_dim)
        self.log_std = nn.Parameter(torch.zeros(1, action_dim))
        self.activ = [nn.ReLU(), nn.Tanh()][use_tanh]
        if use_orthogonal_init:
            orthogonal_init(self.l1)
            orthogonal_init(self.l2)
            orthogonal_init(self.mean_layer, gain=0.01)

    def forward(self, state):
        x = self.activ(self.l1(state))
        x = self.activ(self.l2(x))
        mean = torch.tanh(self.mean_layer(x)) * self.max_action
        return mean

    def get_dist(self, state):
        mean = self.forward(state)
        log_std = self.log_std.expand_as(mean)
        std = torch.exp(log_std)
        dist = torch.distributions.Normal(mean, std)
        return dist
